import os
import logging
from llama_index.tools.tavily_research import TavilyToolSpec
from llama_index.core.schema import Document

# 获取日志记录器
logger = logging.getLogger(__name__)

def fetch_and_save_documents(query: str, max_results: int, output_dir: str, api_key: str) -> None:
    """
    使用 Tavily API 搜索查询，并将结果保存为本地文本文件。

    Args:
        query (str): 要搜索的查询字符串。
        max_results (int): 要获取的最大结果数。
        output_dir (str): 保存文档的目录。
        api_key (str): Tavily API 密钥。
    """
    logger.info(f"开始使用 Tavily 搜索，查询: '{query}'...")
    try:
        tavily_tool = TavilyToolSpec(api_key=api_key)
        documents: list[Document] = tavily_tool.search(query, max_results=max_results)
    except Exception as e:
        logger.error(f"Tavily API 调用失败: {e}")
        return

    if not documents:
        logger.warning("Tavily 未返回任何文档，无需创建文件。")
        return

    os.makedirs(output_dir, exist_ok=True)
    logger.info(f"将在 '{output_dir}' 目录中保存 {len(documents)} 篇文档。")

    for i, doc in enumerate(documents):
        content = doc.text
        filename = f"document_{i + 1}.txt"
        file_path = os.path.join(output_dir, filename)

        try:
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(content)
            logger.debug(f"成功从 ID '{doc.id_}' 提取内容并保存到: {file_path}")
        except IOError as e:
            logger.error(f"写入文件 {file_path} 时出错: {e}")

    logger.info(f"所有文档已成功提取到 '{output_dir}' 目录。")